package com.bruce.ai.alibaba.workflow.config;

import java.util.HashMap;

import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.KeyStrategyFactory;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.agent.ReactAgent;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;


@Configuration
public class MutilAgentWorkflowConfig {

    @Bean
    public StateGraph mutilAgentWorkflow(
        @Qualifier("analysisAgent") ReactAgent analysisAgent,
        @Qualifier("summaryAgent") ReactAgent summaryAgent,
        @Qualifier("researchAgent") ReactAgent researchAgent) throws Exception{
        //定义策略路线
        KeyStrategyFactory keyStrategyFactory = () -> {
            HashMap<String, KeyStrategy> strategies = new HashMap<>();
            strategies.put("input", new ReplaceStrategy());
            strategies.put("research_data", new ReplaceStrategy());
            strategies.put("analysis_result", new ReplaceStrategy());
            strategies.put("final_summary", new ReplaceStrategy());
            return strategies;
        };
        
        // 构建 StateGraph
        StateGraph graph = new StateGraph(keyStrategyFactory);

        /**
         * 定义agent节点，addNode方法中的第一个参数为节点名称，必须和agent的name属性一致
         */
        graph.addNode("research", researchAgent.asNode(false, true, "research_data"));
        graph.addNode("data_analyzer", analysisAgent.asNode(false, true, "analysis_result"));
        graph.addNode("summary", summaryAgent.asNode(false, true, "final_summary"));


        // 定义流程
        graph.addEdge(StateGraph.START, "research");
        graph.addEdge("research", "data_analyzer");
        graph.addEdge("data_analyzer", "summary");
        graph.addEdge("summary", StateGraph.END);

        return graph;
    }

}
